#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/9/14 6:23
@Author   : liujingmao
@File     : vector_database_service.py
"""

import weaviate
from injector import inject
from langchain_openai import OpenAIEmbeddings
from langchain_weaviate import WeaviateVectorStore
from weaviate import WeaviateClient
from weaviate.auth import AuthApiKey


@inject
class VectorDatabaseService:
    """
    向量数据库服务
    """
    client: WeaviateClient
    vector_store: WeaviateVectorStore

    def __init__(self):
        """构造函数，完成向量数据库服务的客户端+LangChain向量数据库实例的创建"""
        # 1. 创建/连接weaviate向量数据库
        self.client = weaviate.connect_to_wcs(
            cluster_url="https://u5elrglqdgrd8ohzixxzq.c0.us-west3.gcp.weaviate.cloud",
            auth_credentials=AuthApiKey("1yebhV73vCkxE8xqfqA2Q6zoI7YzIl3A1RZP"),
        )

        # 2. 创建LangChain向量数据库
        self.vector_store = WeaviateVectorStore(
            client=self.client,
            index_name="DataSet",
            embedding=OpenAIEmbeddings(model="text-embedding-3-small")
        )
